﻿// Copyright (c) 2009, Bo Zhou<http://jedimaster.cnblogs.com>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//     * Neither the name of the organization to which the authors belong nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY AUTHORS AND CONTRIBUTORS ''AS IS'' AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL AUTHORS AND CONTRIBUTORS BE LIABLE FOR ANY
// DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

/*
 *----------------------------------------------------------------------------
 * dja - I've tried not to deviate too far from Bo's original code.
 *----------------------------------------------------------------------------
 */


#ifdef _MSC_VER
// ignore deprecated warnings.. there are quite a few
#pragma warning(disable:4996)
#endif

#include <memory.h>
#include <zlib.h>

#include "fileio.h"
#include "PRT.h"
#include <boost/scoped_array.hpp>

#include "Half/half.cpp"

using namespace std;
using namespace HDK_PRT;

struct Header
{
    char MagicNumber[8];
    int Len;
    char Signature[32];
    int Version;
    long long Count;
};

struct Channel
{
    char Name[32];
    int DataType;
    int Arity;
    int Offset;
};

struct ItoS
{
    int i;
    char *s;
};

#define STRINGIFY(x) #x
#define I2SENTRY(x) {x, (char*)STRINGIFY(x)}

const char*
typeStr(DataType dt)
{
    ItoS i2s[] =
    {
        I2SENTRY(kINT32),
        I2SENTRY(kINT64),
        I2SENTRY(kFLOAT16),
        I2SENTRY(kFLOAT32),
        I2SENTRY(kFLOAT64)
    };
    for(unsigned int i=0;i<sizeof(i2s)/sizeof(i2s[0]);i++)
    {
        if(i2s[i].i == dt)
        {
            return i2s[i].s;
        }
    }
    return "<UNDEFINED>";
}

class TPrintClosure
{
public:
    TPrintClosure(){}
    virtual double print() = 0;
};

template<typename T>
class TPrintAttribDataClosure : public TPrintClosure
{
private:
    double (*mFN)(T, int, int);
    int mArity;
    T *mIterator;
public:
    TPrintAttribDataClosure(double (*fn)(T, int, int), int arity, void *pArray):
      mFN(fn), mArity(arity), mIterator((T*)pArray), TPrintClosure()
    {
    }
    virtual double print()
    {
        double t = 0.0;
        for(int i=0;i<mArity;i++)
        {
            t += (*mFN)(*mIterator, i, mArity);
            mIterator++;
        }
        return t;
    }
};

template<typename T>
double printDataF(T val, int index, int arity)
{
    printf("%f%c", float(val), (index < (arity - 1)) ? ',':' ');
    return double(val);
}

template<typename T>
double printDataI(T val, int index, int arity)
{
    printf("%d%c", int(val), (index < (arity - 1)) ? ',':' ');
    return double(val);
}

class CPrintDataClosure
{
private:
    TPrintClosure *mClosure;
public:
    CPrintDataClosure(DataType type, int arity, void *pArray, long long pnt):
      mClosure(NULL)
    {
        long long i = pnt*arity;
        switch(type)
        {
            case kFLOAT16:
                mClosure = new TPrintAttribDataClosure<half>(printDataF<half>, arity, (unsigned char*)pArray+sizeof(half)*i);
                break;
            case kFLOAT32:
                mClosure = new TPrintAttribDataClosure<float>(printDataF<float>, arity, (unsigned char*)pArray+sizeof(float)*i);
                break;
            case kFLOAT64:
                mClosure = new TPrintAttribDataClosure<double>(printDataF<double>, arity, (unsigned char*)pArray+sizeof(double)*i);
                break;
            case kINT32:
                mClosure = new TPrintAttribDataClosure<int>(printDataI<int>, arity, (unsigned char*)pArray+sizeof(int)*i);
                break;
            case kINT64:
                mClosure = new TPrintAttribDataClosure<long long>(printDataI<long long>, arity, (unsigned char*)pArray+sizeof(long long)*i);
                break;
            default:
                throw std::runtime_error("attribute type is invalid");
        }
    }
    ~CPrintDataClosure()
    {
        if(mClosure)
        {
            delete(mClosure);
        }
    }
    double print()
    {
        return mClosure->print();
    }
};

double
PRT::printValues(bool verbose)
{
    PRTChannelDict dict;

    double total = 0;

    getChannels(dict);
    PRTChannelDict::iterator dictIter;
    long long chCount = getCount();

    if(verbose)
    {
        printf("\nChannels");
        for (dictIter=dict.begin(); dictIter != dict.end(); dictIter++)
        {
            PRTChannel ch = (*dictIter).second;
            printf("\n%u '%s' arity: %d type: %s", (*dictIter).first, ch.mName.c_str(), ch.mArity, typeStr(ch.mType));
        }
    }

    // read in the points
    if(verbose)
    {
        printf("\nPrinting %lld points", chCount);
    }
    for(long long pnt=0;pnt<chCount;pnt++)
    {
        if(verbose)
        {
            printf("\n%llu: ", pnt);
        }
        for (dictIter=dict.begin(); dictIter != dict.end(); dictIter++)
        {
            PRTChannel ch = (*dictIter).second;

            CPrintDataClosure closure(ch.mType, ch.mArity, ch.mData.get(), pnt);
            total += closure.print();
        }
    }
    if(verbose)
    {
        printf("\n");
    }
    return total;
}

inline size_t
getDataLength(const DataType& DT)
{
    switch(DT)
    {
        case kFLOAT16:
        {
            return 2;
        }

        case kINT32:
        case kFLOAT32:
        {
            return 4;
        }

        case kINT64:
        case kFLOAT64:
        {
            return 8;
        }
        default:
        {
            fprintf(stderr, "\nERROR: Unknown data type in getDataLength");
        }
    }
    return 0;
}

void
PRT::addChannel(unsigned int id, const std::string& Attr, const DataType DT, const int Arity, const Array& Data)
{
    if( Attr == "" || DT == kUNKNOWN || Arity == 0 || Data.get() == NULL )
    {
        throw std::runtime_error("cannot add invalid channel");
    }

    PRTChannel Ch;
    Ch.mName = Attr;
    Ch.mArity = Arity;
    Ch.mType = DT;
    Ch.mData = Data;

    mCh.insert( std::make_pair(id, Ch) );
}

void
PRT::saveToFile(long long count, const char* Path)
{
    mCount = count;

    //Create file to write
    FileIO file(Path, "wb");

    //Initialize the header
    Header H;
    memset(&H,0,sizeof(Header));
    H.MagicNumber[0] = (char)192;
    H.MagicNumber[1] = 'P';
    H.MagicNumber[2] = 'R';
    H.MagicNumber[3] = 'T';
    H.MagicNumber[4] = '\r';
    H.MagicNumber[5] = '\n';
    H.MagicNumber[6] = 26;
    H.MagicNumber[7] = '\n';
    H.Len = 56;
    sprintf(H.Signature,"Extensible Particle Format");
    H.Version = 1;
    H.Count = mCount;

    file.write(&H,sizeof(Header),1);

    int RESERVED = 4;
    file.write(&RESERVED,sizeof(int),1);
    file.flush();

    //We need double-sized memory space.
    size_t TotalSize = 0;
    size_t SingleSize = 0; // the size in bytes of all channels for a single data record
    for( std::map<unsigned int, PRTChannel>::iterator itr = mCh.begin(); itr != mCh.end(); itr++)
    {
        const DataType CurrDT = itr->second.mType;
        SingleSize += getDataLength(CurrDT) * itr->second.mArity;
    }
    TotalSize = SingleSize*H.Count;

    boost::shared_array<Bytef> Cache( new Bytef[SingleSize] );
    boost::shared_array<Bytef> OutputBuf( new Bytef[TotalSize*2] );

    // iterate through the channels, copying the data for each to the compression buffer
    for( long long i=0; i<H.Count; i++ )
    {
        size_t Offset = 0;
        for( std::map<unsigned int, PRTChannel>::iterator itr = mCh.begin(); itr != mCh.end(); itr++)
        {
            const DataType CurrDT = itr->second.mType;
            const unsigned char* p = itr->second.mData.get();

            size_t Len = getDataLength(CurrDT) * itr->second.mArity;

            memcpy( Cache.get() + Offset, p + i*Len, Len );

            Offset += Len;
        }

        memcpy( OutputBuf.get()+i*SingleSize, Cache.get(), SingleSize );
    }

    // compress the data
    uLongf ZSrcBufLen = (uLongf)TotalSize;
    Bytef* ZSrcBuf = (Bytef*)OutputBuf.get();

    uLongf ZDstBufLen = compressBound(ZSrcBufLen);
    Bytef* ZDstBuf = (Bytef*)malloc(ZDstBufLen);

    int ZR = compress2(ZDstBuf,&ZDstBufLen,ZSrcBuf,ZSrcBufLen,6);
    if( ZR != Z_OK )
    {
        throw std::runtime_error("compress2 failure");
    }

    // write channel metadata
    int ChanNumer = (int)mCh.size();
    file.write(&ChanNumer,sizeof(int),1);
    int ChanLen = 44;
    file.write(&ChanLen,sizeof(int),1);
    file.flush();

    int Offset = 0;
    for( std::map<unsigned int, PRTChannel>::iterator itr = mCh.begin(); itr != mCh.end(); itr++)
    {
        Channel Ch;
        memset(&Ch,0,sizeof(Channel));

        Ch.Arity = itr->second.mArity;
        Ch.DataType = (int)itr->second.mType;
        strncpy(Ch.Name, itr->second.mName.c_str(), 31);
        Ch.Offset = Offset;

        const DataType CurrDT = itr->second.mType;

        Offset += (int)getDataLength(CurrDT) * Ch.Arity;

        file.write(&Ch,sizeof(Channel),1);
    }

    // write the compressed data
    file.write(ZDstBuf,ZDstBufLen,1);
    delete(ZDstBuf);
}

void
PRT::readFromFile(const char * Path)
{

    mCount = 0;

    FileIO file(Path,"rb");

    Header H;
    memset(&H,0,sizeof(Header));

    size_t result = file.read(&H,sizeof(Header),1);

    if( strncmp(H.Signature,"Extensible Particle Format",31) != 0 )
    {
        throw std::runtime_error("invalid signature");
    }

    if( H.Len != 56  )
    {
        throw std::runtime_error("invalid len");
    }

    if( H.Version != 1 )
    {
        throw std::runtime_error("invalid version");
    }

    int REVERSED = 0;
    result = file.read(&REVERSED,sizeof(int),1);

    if( REVERSED != 4 )
    {
        throw std::runtime_error("invalid REVERSED value");
    }

    int ChDef[2] = {0,0};
    result = file.read(&ChDef[0],sizeof(int)*2,1);
    if( ChDef[0] == 0 || ChDef[1] != 44 )
    {
        throw std::runtime_error("invalid channel data");
    }

    const int ChCount = ChDef[0];

    boost::scoped_array<Channel> Channels(new Channel[ChCount]);
    memset(Channels.get(),0,sizeof(Channel)*ChCount);
    result = file.read(Channels.get(),sizeof(Channel),ChCount);

    size_t SingleSize = 0;
    unsigned int id = 0;
    for( long long i=0; i<ChDef[0]; i++ )
    {
        DataType CurrDT = (DataType)Channels[i].DataType;
        PRTChannel Ch;
        Ch.mName = Channels[i].Name;
        Ch.mType = CurrDT;
        Ch.mArity = Channels[i].Arity;
        size_t Len = getDataLength(CurrDT)*Ch.mArity;

        Ch.mData.reset(new unsigned char[H.Count*Len]);

        mCh.insert( make_pair(id,Ch) );
        id++;
        SingleSize += Len;
    }

    if( SingleSize == 0 )
    {
        throw std::runtime_error("total channel length is zero");
    }

    long CurrFilePos = file.tell();
    file.seek(0, SEEK_END);
    unsigned long ZSrcLen = file.tell() - CurrFilePos;
    file.seek(CurrFilePos, SEEK_SET);

    boost::scoped_array<Bytef> ZSrc(new Bytef[ZSrcLen]);
    result = file.read(ZSrc.get(),ZSrcLen,1);

    uLongf ZDstLen = (uLongf)(SingleSize*H.Count);
    boost::scoped_array<Bytef> ZDst(new Bytef[ZDstLen]);

    int ZE = uncompress(ZDst.get(),&ZDstLen,ZSrc.get(),ZSrcLen);

    if( ZE != Z_OK )
    {
        throw std::runtime_error("uncompress failed");
    }

    int AccOffset = 0;
    int mChIdx = 0;
    unsigned char* WholePointer = (unsigned char*)ZDst.get();
    for( PRTChannelDict::iterator itr = mCh.begin(); itr != mCh.end(); itr++ )
    {
        DataType CurrDT = (DataType)itr->second.mType;

        int Size = (int)getDataLength(CurrDT)*itr->second.mArity;

        unsigned char* RawPointer = itr->second.mData.get();
        for( long long int n=0; n<H.Count; n++ )
        {
            memcpy( RawPointer+n*Size, WholePointer+SingleSize*n+AccOffset, Size );
        }
        AccOffset += Size;
        mChIdx++;
    }

    mCount = H.Count;
}

void
PRT::clear()
{
    mCount = 0;
    mCh.clear();
}

void
PRT::getChannels(PRTChannelDict& Dict) const
{
    Dict = mCh;
}

long long
PRT::getCount() const
{
    return mCount;
}
